from mcp.server.fastmcp import FastMCP
import secrets

mcp = FastMCP("Large Prime Generator Server")

def is_prime(n: int, k: int = 5) -> bool:
    """Miller-Rabin素性测试"""
    if n <= 1:
        return False
    elif n <= 3:
        return True
    elif n % 2 == 0:
        return False
    
    d = n - 1
    s = 0
    while d % 2 == 0:
        d //= 2
        s += 1
    
    for _ in range(k):
        a = secrets.randbelow(n - 2) + 2
        x = pow(a, d, n)
        if x == 1 or x == n - 1:
            continue
        for __ in range(s - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                break
        else:
            return False
    return True

@mcp.tool()
def generate_large_prime(bits: int = 1024) -> int:
    """
    生成指定位数的大素数
    :param bits: 素数位数（默认1024位）
    :return: 大素数
    """
    while True:
        candidate = secrets.randbits(bits)
        candidate |= (1 << (bits - 1)) | 1  # 确保最高位和最低位为1
        if is_prime(candidate):
            return candidate

if __name__ == "__main__":
    mcp.run()